package main

import (
	"encoding/base64"
	"fmt"
	"log"
	"net/http"
	"os/exec"
	"strings"

	"github.com/gorilla/websocket"
)

type RTSPStreamer struct {
	InputURL  string
	OutputDir string
	Process   *exec.Cmd
}

// 添加 WebSocket 升级器
var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许所有跨域请求
	},
}

func NewRTSPStreamer(inputURL, outputDir string) *RTSPStreamer {
	return &RTSPStreamer{
		InputURL:  inputURL,
		OutputDir: outputDir,
	}
}

// 原有的 HLS 流转换
func (s *RTSPStreamer) StartStreaming() error {
	args := []string{
		"-i", s.InputURL,
		"-fflags", "nobuffer",
		"-rtsp_transport", "tcp",
		//"-c:v", "copy",
		//"-c:a", "aac",
		//"-f", "hls",
		"-hls_time", "10",
		"-hls_list_size", "5",
		"-hls_flags", "delete_segments",
		fmt.Sprintf("%s/stream.m3u8", s.OutputDir),
	}

	s.Process = exec.Command("ffmpeg", args...)
	stderr, err := s.Process.StderrPipe()
	if err != nil {
		return fmt.Errorf("error getting stderr pipe: %v", err)
	}

	if err := s.Process.Start(); err != nil {
		return fmt.Errorf("error starting ffmpeg: %v", err)
	}

	go func() {
		buf := make([]byte, 2048)
		for {
			n, err := stderr.Read(buf)
			if n > 0 {
				log.Printf("ffmpeg: %s", strings.TrimSpace(string(buf[:n])))
			}
			if err != nil {
				break
			}
		}
	}()

	return nil
}

// 新增 FLV 流处理函数
func handleRTSPToFLV(w http.ResponseWriter, r *http.Request) {
	log.Println("开始处理新的 RTSP 转 FLV 请求")

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket 升级失败: %v", err)
		return
	}
	defer conn.Close()
	log.Println("WebSocket 连接已建立")

	// 获取并解码 base64 编码的 RTSP URL
	encodedURL := r.URL.Query().Get("url")
	if encodedURL == "" {
		log.Println("未提供 RTSP URL")
		return
	}
	log.Println("已获取编码后的 URL")

	// Base64 解码
	decodedBytes, err := base64.StdEncoding.DecodeString(encodedURL)
	if err != nil {
		log.Printf("Base64 解码失败: %v", err)
		return
	}
	rtspURL := string(decodedBytes)
	log.Printf("解码后的 RTSP URL: %s", rtspURL)

	// 修复 FFMPEG 参数 - 取消注释关键参数
	args := []string{
		"-i", rtspURL,
		"-rtsp_transport", "tcp",
		"-buffer_size", "204800",
		"-c:v", "libx264",
		"-threads", "8",
		"-tune", "zerolatency",
		"-preset", "ultrafast",
		"-b:v", "4M",
		"-maxrate", "4M",
		"-bufsize", "8M",
		"-r", "25",
		"-an",
		"-f", "flv",
		"pipe:1",
	}
	log.Printf("FFMPEG 命令参数: %v", args)

	cmd := exec.Command("ffmpeg", args...)

	// 添加错误输出管道以便调试
	stderr, err := cmd.StderrPipe()
	if err != nil {
		log.Printf("创建错误输出管道失败: %v", err)
		return
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Printf("创建输出管道失败: %v", err)
		return
	}
	log.Println("输出管道创建成功")

	if err := cmd.Start(); err != nil {
		log.Printf("启动 FFMPEG 失败: %v", err)
		return
	}
	log.Println("FFMPEG 进程已启动")

	// 读取并打印 FFMPEG 错误输出
	go func() {
		buf := make([]byte, 2048)
		for {
			n, err := stderr.Read(buf)
			if n > 0 {
				log.Printf("FFMPEG 错误输出: %s", strings.TrimSpace(string(buf[:n])))
			}
			if err != nil {
				log.Printf("FFMPEG 错误输出结束: %v", err)
				break
			}
		}
	}()

	go func() {
		defer cmd.Process.Kill()
		log.Println("开始传输视频流")
		buf := make([]byte, 8192)
		var dataCount int64 = 0
		var frameCount int = 0

		for {
			n, err := stdout.Read(buf)
			if n > 0 {
				dataCount += int64(n)
				frameCount++

				// 每100帧打印一次统计信息
				if frameCount%100 == 0 {
					log.Printf("已传输 %d 帧数据，总计 %.2f KB", frameCount, float64(dataCount)/1024)
				}

				if err := conn.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil {
					log.Printf("WebSocket 写入错误: %v", err)
					return
				}
			}
			if err != nil {
				log.Printf("读取流结束或发生错误: %v，共传输 %d 帧，%.2f KB 数据",
					err, frameCount, float64(dataCount)/1024)
				break
			}
		}
		log.Println("视频流传输结束")
	}()

	log.Println("等待 WebSocket 连接关闭")
	for {
		if _, _, err := conn.ReadMessage(); err != nil {
			log.Printf("WebSocket 连接已关闭: %v", err)
			break
		}
	}
	log.Println("RTSP 转 FLV 处理完成")
}

func (s *RTSPStreamer) StopStreaming() error {
	if s.Process != nil && s.Process.Process != nil {
		return s.Process.Process.Kill()
	}
	return nil
}
